# -*- coding: utf-8 -*-
"""
@Time    : 2024/6/27 11:27 
@Author  : ZhangShenao 
@File    : 1.手写简易Chain.py 
@Desc    : 手写简易Chain
"""
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

from simple_chain import Chain

# 编排PromptTemplate
chat_prompt_template = ChatPromptTemplate.from_messages(
    [
        ('system', '你是智普AI开发的聊天机器人，请根据用户的提问进行回答。'),
        ('human', '{query}')
    ]
)

# 创建智普AI LLM
llm = ChatZhipuAI(
    model="GLM-4-Air",
    temperature=0.5,
)

# 创建OutputParser
output_parser = StrOutputParser()

# 创建执行链Chain,用于编排执行流程
chain = Chain(steps=[chat_prompt_template, llm, output_parser])

# 执行链,获取结果
content = chain.invoke({
    'query': '请帮我讲一个关于程序员的冷笑话',
})

print(content)
